#include <jni.h>
#include <string>
#include <iostream>
#include "include/x264.h"
#include <android/log.h>

#include<android/log.h>
#define LOG_TAG "System.out"//System.out可以修改为其他字符串，该字符串会作为输出日志的TAG。
#define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG,LOG_TAG,__VA_ARGS__)
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO,LOG_TAG,__VA_ARGS__)


/**
 * x264编码测试函数,测试x264是否可以正常使用
 */
extern "C" JNIEXPORT jboolean JNICALL
Java_com_dream_libx264_XH264Encoder_test(
        JNIEnv* env,
        jobject /* this */) {
    // 初始化x264编码器参数
    x264_param_t param;
    x264_param_default(&param);

    // 设置编码器参数
    param.i_width = 640;  // 视频宽度
    param.i_height = 480; // 视频高度
    param.i_fps_num = 25; // 帧率分子
    param.i_fps_den = 1;  // 帧率分母

    // 配置x264编码器
    x264_param_apply_profile(&param, "baseline");

    // 打开x264编码器
    x264_t* encoder = x264_encoder_open(&param);
    if (!encoder) {
        LOGD("Failed to open x264 encoder.");
        return false;
    }

    // 创建输入图像帧
    x264_picture_t pic_in;
    x264_picture_alloc(&pic_in, X264_CSP_I420, param.i_width, param.i_height);

    // 填充虚拟数据（这里简单地将所有像素置为0）
    for (int i = 0; i < param.i_height; ++i) {
        for (int j = 0; j < param.i_width; ++j) {
            pic_in.img.plane[0][i * param.i_width + j] = 0; // Y分量
        }
    }
    for (int i = 0; i < param.i_height / 2; ++i) {
        for (int j = 0; j < param.i_width / 2; ++j) {
            pic_in.img.plane[1][i * (param.i_width / 2) + j] = 128; // U分量
            pic_in.img.plane[2][i * (param.i_width / 2) + j] = 128; // V分量
        }
    }

    // 编码一帧数据
    x264_nal_t* nals;
    int i_nal;
    x264_picture_t pic_out;
    int frame_size = x264_encoder_encode(encoder, &nals, &i_nal, &pic_in, &pic_out);
    if (frame_size < 0) {
        LOGD("Failed to encode frame.");
        x264_picture_clean(&pic_in);
        x264_encoder_close(encoder);
        return false;
    }

    LOGD("Frame encoded successfully. NAL units: %d, Frame size: %d",i_nal,frame_size);

    // 清理资源
    x264_picture_clean(&pic_in);
    x264_encoder_close(encoder);
    return true;
}

/**
 * x264编码测试函数2,将YUV文件转为H264文件
 */
extern "C"
JNIEXPORT jboolean JNICALL
Java_com_dream_libx264_XH264Encoder_test2(JNIEnv *env, jobject thiz) {


    FILE *fp_src = fopen("/mnt/sdcard/cwt/1738908338941.yuv", "rb");
    FILE *fp_dst = fopen("/mnt/sdcard/cwt/1738908338941.h264", "wb");
    if (fp_src == NULL || fp_dst == NULL) {
        LOGD("openErr");
        return false;
    }
    x264_picture_t *pPic_in = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    x264_picture_t *pPic_out = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    x264_param_t *pParam = (x264_param_t *) (malloc(sizeof(x264_param_t)));
    x264_param_default(pParam);
    pParam->i_width = 1920;
    pParam->i_height = 1080;
    int csp = X264_CSP_I420;
    pParam->i_csp = csp;
    x264_param_apply_profile(pParam, x264_profile_names[5]);
    int iNal = 0;
    x264_nal_t *pNals = NULL;
    x264_t *pHandle = x264_encoder_open_164(pParam);
    LOGD("pHandle:%d", pHandle);
    x264_picture_init(pPic_out);
    x264_picture_alloc(pPic_in, csp, pParam->i_width, pParam->i_height);
    int y_size = pParam->i_width * pParam->i_height;

    //计算YUV帧数
    fseek(fp_src, 0, SEEK_END);
    long fileSize = ftell(fp_src);
    rewind(fp_src);
    int frameSize = pParam->i_width * pParam->i_height * 3 / 2;
    int numFrames = (int)(fileSize / frameSize);
    LOGD("帧数:%d", numFrames);

    for (int i = 0; i < numFrames; ++i) {
        fread(pPic_in->img.plane[0], y_size, 1, fp_src);         //Y
        fread(pPic_in->img.plane[1], y_size / 4, 1, fp_src);     //U
        fread(pPic_in->img.plane[2], y_size / 4, 1, fp_src);     //V

        pPic_in->i_pts = i;
        int ret = x264_encoder_encode(pHandle, &pNals, &iNal, pPic_in, pPic_out);
        LOGD("x264_encoder_encode:%d", ret);
        LOGD("iNal:%d", iNal);
        if (ret < 0) {
            return false;;
        }

        for (int j = 0; j < iNal; ++j) {
            LOGD("type:%d", pNals[j].i_type);
            fwrite(pNals[j].p_payload, 1, pNals[j].i_payload, fp_dst);
        }
    }

    //flush encoder
    while (1) {
        int ret = x264_encoder_encode(pHandle, &pNals, &iNal, NULL, pPic_out);
        if (ret == 0) {
            break;
        }
        printf("Flush 1 frame.\n");
        for (int j = 0; j < iNal; ++j) {
            fwrite(pNals[j].p_payload, 1, pNals[j].i_payload, fp_dst);
        }
    }
    x264_picture_clean(pPic_in);
    x264_encoder_close(pHandle);
    pHandle = NULL;

    free(pPic_in);
    free(pPic_out);
    free(pParam);

    fclose(fp_src);
    fclose(fp_dst);
    return true;
}

struct EncoderContext {
    x264_t *pHandle;
    x264_picture_t *pPic_in;
    x264_picture_t *pPic_out;
    x264_param_t *pParam;
    JNIEnv *env;
    jobject obj;
    jmethodID callbackFrameMethod,callbackFormatMethod;
};

extern "C" JNIEXPORT jlong JNICALL
Java_com_dream_libx264_XH264Encoder_initEncoder(JNIEnv *env, jobject thiz, jint width, jint height) {
    EncoderContext *ctx = new EncoderContext();
    ctx->pParam = (x264_param_t *) malloc(sizeof(x264_param_t));

    // 配置x264编码器,第二个参数为编码速度，越小越快,第三个参数为画质,越小画质越好
    x264_param_default_preset(ctx->pParam, x264_preset_names[5], x264_tune_names[5]);

    // 配置x264编码器级别，如baseline，main，high等
    x264_param_apply_profile(ctx->pParam, x264_profile_names[2]);
    ctx->pParam->i_level_idc = 41;

    ctx->pParam->i_width = width;
    ctx->pParam->i_height = height;
    ctx->pParam->i_fps_num = 15;//帧率
    ctx->pParam->i_fps_den = 1;//帧率分母

    ctx->pParam->rc.i_rc_method = X264_RC_CRF;
    ctx->pParam->rc.f_rf_constant = 23.0f;

    //设置运动估计和子像素运动估计
    ctx->pParam->analyse.i_me_method = X264_ME_UMH;
    ctx->pParam->analyse.i_subpel_refine = 7;

    ctx->pParam->i_frame_reference = 3;
    ctx->pParam->i_bframe = 0;//不使用B帧
    ctx->pParam->i_bframe_pyramid = 0;

    ctx->pParam->analyse.i_me_method = X264_ME_HEX;
    ctx->pParam->analyse.i_me_range = 16; // Adjust motion estimation range
    
    ctx->pParam->analyse.f_psy_rd = 1.0f;
    // Adaptive quantization
    ctx->pParam->rc.i_aq_mode = 3;

    // 设置去块滤波器参数
    ctx->pParam->b_deblocking_filter = 1;
    ctx->pParam->i_deblocking_filter_alphac0 = -3;
    ctx->pParam->i_deblocking_filter_beta = -3;

    // Threads and lookahead
    ctx->pParam->i_threads = X264_THREADS_AUTO;
    ctx->pParam->i_sync_lookahead = 9; // Set rc-lookahead to 9

    // QP max
    ctx->pParam->rc.i_qp_max = 40;
    ctx->pParam->rc.i_qp_min = 10;

    //关键帧间隔
    ctx->pParam->i_keyint_max = 125;
    ctx->pParam->i_keyint_min = 15;

    ctx->pHandle = x264_encoder_open_164(ctx->pParam);
    if (!ctx->pHandle) {
        LOGD("Failed to open encoder");
        delete ctx;
        return 0;
    }

    ctx->pPic_in = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    ctx->pPic_out = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    x264_picture_init(ctx->pPic_out);
    x264_picture_alloc(ctx->pPic_in, ctx->pParam->i_csp, ctx->pParam->i_width, ctx->pParam->i_height);

    jclass cls = env->GetObjectClass(thiz);
    ctx->callbackFrameMethod = env->GetMethodID(cls, "onEncodedDataAvailableCallback", "([BJI)V");

    if (ctx->callbackFrameMethod == nullptr) {
        LOGD("Failed to get method ID for onEncodedDataAvailableCallback");
        x264_encoder_close(ctx->pHandle);
        free(ctx->pPic_in);
        free(ctx->pPic_out);
        free(ctx->pParam);
        delete ctx;
        return 0;
    }

    ctx->obj = env->NewGlobalRef(thiz);
    ctx->env = env;

    return reinterpret_cast<jlong>(ctx);
}

extern "C" JNIEXPORT void JNICALL
Java_com_dream_libx264_XH264Encoder_encode(JNIEnv *env, jobject thiz, jlong handle, jbyteArray yuvData, jlong pts) {
    EncoderContext *ctx = reinterpret_cast<EncoderContext *>(handle);
    jsize length = env->GetArrayLength(yuvData);
    jbyte *yuvBytes = env->GetByteArrayElements(yuvData, nullptr);

    if (length != ctx->pParam->i_width * ctx->pParam->i_height * 3 / 2) {
        LOGD("Invalid frame size");
        env->ReleaseByteArrayElements(yuvData, yuvBytes, 0);
        return;
    }

    memcpy(ctx->pPic_in->img.plane[0], yuvBytes, ctx->pParam->i_width * ctx->pParam->i_height);
    memcpy(ctx->pPic_in->img.plane[1], yuvBytes + ctx->pParam->i_width * ctx->pParam->i_height,
           ctx->pParam->i_width * ctx->pParam->i_height / 4);
    memcpy(ctx->pPic_in->img.plane[2], yuvBytes + ctx->pParam->i_width * ctx->pParam->i_height * 5 / 4,
           ctx->pParam->i_width * ctx->pParam->i_height / 4);

    ctx->pPic_in->i_pts = pts;

    int iNal = 0;
    x264_nal_t *pNals = NULL;
    int ret = x264_encoder_encode(ctx->pHandle, &pNals, &iNal, ctx->pPic_in, ctx->pPic_out);
    if (ret < 0) {
        LOGD("Encoding failed");
    } else {
        for (int j = 0; j < iNal; ++j) {
            jbyteArray encodedData = env->NewByteArray(pNals[j].i_payload);
            env->SetByteArrayRegion(encodedData, 0, pNals[j].i_payload, reinterpret_cast<const jbyte*>(pNals[j].p_payload));
            env->CallVoidMethod(ctx->obj, ctx->callbackFrameMethod, encodedData, ctx->pPic_out->i_pts,pNals[j].i_type);
            env->DeleteLocalRef(encodedData);
        }
    }

    env->ReleaseByteArrayElements(yuvData, yuvBytes, 0);
}

extern "C" JNIEXPORT void JNICALL
Java_com_dream_libx264_XH264Encoder_flushAndCloseEncoder(JNIEnv *env, jobject thiz, jlong handle) {
    EncoderContext *ctx = reinterpret_cast<EncoderContext *>(handle);

    int iNal = 0;
    x264_nal_t *pNals = NULL;
    while (1) {
        int ret = x264_encoder_encode(ctx->pHandle, &pNals, &iNal, NULL, ctx->pPic_out);
        if (ret == 0) {
            break;
        }
        for (int j = 0; j < iNal; ++j) {
            bool isIFrame = (pNals[j].i_type == NAL_SLICE_IDR);
            jbyteArray encodedData = env->NewByteArray(pNals[j].i_payload);
            env->SetByteArrayRegion(encodedData, 0, pNals[j].i_payload, reinterpret_cast<const jbyte*>(pNals[j].p_payload));
            env->CallVoidMethod(ctx->obj, ctx->callbackFrameMethod, encodedData, ctx->pPic_out->i_pts, isIFrame ? JNI_TRUE : JNI_FALSE);
            env->DeleteLocalRef(encodedData);
        }
    }

    x264_picture_clean(ctx->pPic_in);
    x264_encoder_close(ctx->pHandle);

    env->DeleteGlobalRef(ctx->obj);
    free(ctx->pPic_in);
    free(ctx->pPic_out);
    free(ctx->pParam);
    delete ctx;
}